#include "segment.h"

#include <queue>
#include <set>
#include <map>

#include <lemon/list_graph.h>
#include <lemon/bin_heap.h>
#include <lemon/unionfind.h>

using namespace lemon;
using namespace std;
using namespace sds;

int ReadSegments(const SDataView<Point>& src,
                 const SDataView<unsigned char>& seg,
                 SDataView<int>& ren, 
                 std::vector<Cluster>& inf) {
  int num = ReadSegments(seg, ren);

  std::vector<SInteger> size = seg.getSize();

  inf.clear();
  inf.resize(num, Cluster());

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (ren[i][j] != -1) {
        inf[ren[i][j]] += src[i][j];
      }
    }
  }
  return num;
}

int ReadSegments(const SDataView<unsigned char>& seg, 
                 SDataView<int>& ren) {
    
  std::vector<SInteger> size = seg.getSize();

  SDataView<bool> processed(size);
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      processed[i][j] = false;
    }
  }

  int num = 0;
  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      unsigned char current = seg[i][j];
      if (current == 0) {
        ren[i][j] = -1;
      } else if (!processed[i][j]) {
        
        std::queue<Pixel> pixelQueue;
        
        processed[i][j] = true;
        pixelQueue.push(Pixel(i, j));

        while (!pixelQueue.empty()) {
          Pixel pixel = pixelQueue.front();
          pixelQueue.pop();

          ren[pixel.i][pixel.j] = num;
          
          if (pixel.i > 0 && seg[pixel.i - 1][pixel.j] == current
              && !processed[pixel.i - 1][pixel.j]) {
            processed[pixel.i - 1][pixel.j] = true;
            pixelQueue.push(Pixel(pixel.i - 1, pixel.j));            
          }
          if (pixel.j > 0 && seg[pixel.i][pixel.j - 1] == current
              && !processed[pixel.i][pixel.j - 1]) {
            processed[pixel.i][pixel.j - 1] = true;
            pixelQueue.push(Pixel(pixel.i, pixel.j - 1));            
          }
          if (pixel.i + 1 < size[0] && seg[pixel.i + 1][pixel.j] == current
              && !processed[pixel.i + 1][pixel.j]) {
            processed[pixel.i + 1][pixel.j] = true;
            pixelQueue.push(Pixel(pixel.i + 1, pixel.j));            
          }
          if (pixel.j + 1 < size[1] && seg[pixel.i][pixel.j + 1] == current
              && !processed[pixel.i][pixel.j + 1]) {
            processed[pixel.i][pixel.j + 1] = true;
            pixelQueue.push(Pixel(pixel.i, pixel.j + 1));            
          }
        }
        ++num;
      }
    }
  }
  return num;
}

void WriteSegments(const SDataView<int>& ren,
                   SDataView<unsigned char>& seg) {
  vector<SInteger> size = ren.getSize();
  
  typedef ListGraph Graph;

  GRAPH_TYPEDEFS(Graph);

  Graph graph;
  std::map<unsigned, Node> nodes;

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (ren[i][j] == 0) continue;
      if (nodes.find(ren[i][j]) == nodes.end()) {
	Node node = graph.addNode();
	nodes.insert(make_pair(ren[i][j], node));
      }
    }
  }

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (ren[i][j] == 0) continue;
      if (j + 1 < size[1]) {
	if (ren[i][j] != ren[i][j + 1] && ren[i][j + 1] != 0) {
	  if (findEdge(graph, nodes[ren[i][j]], nodes[ren[i][j + 1]]) 
	      == INVALID) {
	    graph.addEdge(nodes[ren[i][j]], nodes[ren[i][j + 1]]);
	  }
	}
      }
      if (i + 1 < size[0]) {
	if (ren[i + 1][j] != 0 && ren[i][j] != ren[i + 1][j]) {
	  if (findEdge(graph, nodes[ren[i][j]], nodes[ren[i + 1][j]]) 
	      == INVALID) {
	    graph.addEdge(nodes[ren[i][j]], nodes[ren[i + 1][j]]);
	  }
	}
      }
    }
  }

  Graph::NodeMap<char> color(graph, -2);
  
  Graph::NodeMap<int> heapMap(graph, -1);
  BinHeap<int, Graph::NodeMap<int> > heap(heapMap);
  
  for (NodeIt it(graph); it != INVALID; ++it) {
    heap.push(it, countOutArcs(graph, it));
  }

  vector<Node> order;

  while (!heap.empty()) {
    Node node = heap.top();
    heap.pop();
    color[node] = -1;
    order.push_back(node);
    for (IncEdgeIt it(graph, node); it != INVALID; ++it) {
      Node target = graph.runningNode(it); 
      if (color[target] == -2) {
	heap.decrease(target, heap[target] - 1);
      }
    }
  }

  for (int i = order.size() - 1; i >= 0; --i) {
    set<int> colors;
    for (IncEdgeIt it(graph, order[i]); it != INVALID; ++it) {
      Node target = graph.runningNode(it); 
      if (color[target] != -1) {
	colors.insert(color[target]);
      }
    }
    for (int j = 1; j <= 6; ++j) {
      if (colors.find(j) == colors.end()) {
	color[order[i]] = j;
	break;
      }
    }
  }

  for (int i = 0; i < size[0]; ++i) {
    for (int j = 0; j < size[1]; ++j) {
      if (ren[i][j] == 0) {
	seg[i][j] = 0;
      } else {
	seg[i][j] = color[nodes[ren[i][j]]];
      }
    }
  }
  
}
